from PyQt4 import QtGui, QtCore, QtOpenGL

import osg2
import lucid.core
from pydispatch import dispatcher


class OpenSGTreeModel(QtCore.QAbstractItemModel):
   """ Base tree model.
   
   @todo: This could probably be refactored a bit now with ideas from property tree.
   """
   lucid_mime_type = "application/lucid"
   
   def __init__(self, parent=None):
      QtCore.QAbstractItemModel.__init__(self, parent)
      self.mRootNode = None
      self.crapTastic = {}    # Temporary map of id to container            
      #pt = osg2.helpers.PrintTreeTraverser()
      #pt.traverse(self.mRootNode)
      dispatcher.connect(self.onNeedLayoutChanged, lucid.core.signal.SelectedFieldContainerChanged)
      
   def getContainer(self, index):
      return self.crapTastic.get(index, None)
   
   def makeContainerId(self, container):
      cont_id = 0
      if container:
         cont_id = osg2.getContainerId(container)
         if not self.crapTastic.has_key(cont_id):
            self.crapTastic[cont_id] = container
      return cont_id
   
   def setRoot(self, rootNode):
      """ Set to a root node.  Reset the tree and data. """
      self.mRootNode = rootNode
      self.crapTastic.clear()
      self.reset()               # Reset the model, triggers view reloads
      
   def flags(self, index):
      if not index.isValid():
         return QtCore.Qt.ItemIsEnabled
      
      col = index.column()
      if 0 == col:
         #return QtCore.Qt.ItemIsEnabled | QtCore.Qt.ItemIsSelectable
         return QtCore.Qt.ItemIsSelectable | QtCore.Qt.ItemIsEnabled | QtCore.Qt.ItemIsEditable \
              | QtCore.Qt.ItemIsDragEnabled | QtCore.Qt.ItemIsDropEnabled
      else:
         return QtCore.Qt.ItemIsEnabled

   def index(self, row, column, parent):
      # I think there may be problems in here
      # Find node and return index for it
      cur_node = None
      parent_node = None
      
      # Check for top level
      if not parent.isValid():  
         if self.mRootNode:            
            return self.createIndex(row,column, self.makeContainerId(self.mRootNode))
         else:
            return QtCore.QModelIndex()     # Return invalid if the tree is empty         
      
      parent_id = parent.internalId()
      parent_node = self.getContainer(parent_id) 
         
      assert not parent_node is None
      assert isinstance(parent_node, osg2.Node)
      num_children = parent_node.getNChildren()
      if row < parent_node.getNChildren():
         cur_node = parent_node.getChild(row)
      else:
         return QtCore.QModelIndex()
            
      assert not cur_node is None      
      return self.createIndex(row, column, self.makeContainerId(cur_node))
   

   def parent(self, child_index):
      # If invalid, then return invalid
      if not child_index.isValid():
         return QtCore.QModelIndex()

      child_id = int(child_index.internalId())      
      child_node = self.getContainer(child_id)
      assert None != child_node
      
      if (child_node == self.mRootNode):
         return QtCore.QModelIndex()
      
      parent_node = child_node.getParent()      
      #if (None == parent_node) or (parent_node == self.mRootNode):
      #   return QtCore.QModelIndex()
      if (None == parent_node) or (parent_node == self.mRootNode):
         return self.createIndex(0,0,self.makeContainerId(self.mRootNode))
      
      # Find parent's index
      grandparent = parent_node.getParent()      
      parent_index = grandparent.findChild(parent_node)
      assert (parent_index >= 0)
            
      return self.createIndex(parent_index, 0, self.makeContainerId(parent_node))

   def rowCount(self, parent):
      if not parent.isValid():
         return 1               # Root node
      else:
         parent_node = self.getContainer(parent.internalId())
         return parent_node.getNChildren()   
   
   def columnCount(self, parent):
      # Just have 1 column for now
      return 2      

   def data(self, index, role):
      if not index.isValid():
         if role == QtCore.Qt.UserRole:
            return None
         else:
            return QtCore.QVariant()
      
      if role in [QtCore.Qt.DisplayRole, 
                  QtCore.Qt.EditRole,
                  QtCore.Qt.TextColorRole,
                  QtCore.Qt.UserRole]:
         node = self.getContainer(index.internalId())
         core = node.getCore()
         col = index.column()

         # Name to display
         if role == QtCore.Qt.DisplayRole or role == QtCore.Qt.EditRole:
            if 0 == col:
               return QtCore.QVariant(str(osg2.getName(node)))
            else:
               type_name = "None"
               if core:
                  type_name = core.getType().getCName()
               return QtCore.QVariant(str(type_name))
         # Userdata reference to node
         elif QtCore.Qt.UserRole == role and 0 == col:
            return node
         # Text color to signify selected nodes.
         elif QtCore.Qt.TextColorRole == role:
            cur_fc = lucid.core.Environment().selectedFC
            
            if isinstance(core, osg2.Geometry):               
               # XXX: Refactor this to color it based on how the fc is
               # referenced               
               if core.getMaterial() == cur_fc:
                  return QtCore.QVariant(QtGui.QColor(QtCore.Qt.darkGreen))
            if cur_fc == core:
               return QtCore.QVariant(QtGui.QColor(QtCore.Qt.darkBlue))
               
      
      # Catch all
      return QtCore.QVariant()
   
   def headerData(self, col, orientation, role):
      if orientation == QtCore.Qt.Horizontal and role == QtCore.Qt.DisplayRole:
         headers = ["Node","Core"]         
         return QtCore.QVariant(str(headers[col]))
      return QtCore.QVariant()


   # ------------ Edit Interface ------------------- #
   def setData(self, index, value, role):
      if not index.isValid():
         return False
      else:
         node = self.getContainer(index.internalId())
         if node and (role == QtCore.Qt.EditRole):
            osg2.setName(node, str(value.toString()))
            self.emit(QtCore.SIGNAL("dataChanged(QModelIndex,QModelIndex)"), index, index)
            self.emit(QtCore.SIGNAL("dataChanged(int)"), index.row())
            return True
      return False

   def supportedDropActions(self):
      # Hold shift when copying to change modes
      # See qdirmodel for a good example or puzzle in pyqt
      print "Supported drop actions called."
      return (QtCore.Qt.CopyAction | QtCore.Qt.MoveAction)
   
   def mimeTypes(self):
      """ List of types we can represent. """
      types = QtCore.QStringList()
      types.append("application/lucid")
      return types
   
   def mimeData(self, indexes):
      index = indexes[0]
      
      if not index.isValid():
         return None
      id = index.internalId()   
      
      mime_data = QtCore.QMimeData()      
      text = "lucid-nodeid:%s"%id
      print "Set mime data: ", text
      mime_data.setData(OpenSGTreeModel.lucid_mime_type,text)
      return mime_data
   
   def dropMimeData(self, mimeData, action, row, column, parent):
      """ Called when we drop a node.
      if row and col == (-1,-1) then just need to parent the node.
      Otherwise, the row is saying which child number we would like to be.
      """
      if not parent.isValid():
         return False
      
      data = str(mimeData.data(OpenSGTreeModel.lucid_mime_type))
      (data_type,id) = data.split(":")
      print "text: ", data
      
      if (action == QtCore.Qt.CopyAction) or (action == QtCore.Qt.MoveAction):
         parent_node = self.data(parent, QtCore.Qt.UserRole)
         cur_node = lucid.core.Environment().selectedNode
         parent_name = "None"
         if parent_node:     
            parent_name = osg2.getName(parent_node)
         print "row: %s  col: %s  parent: %s"%(row,column,parent_name)

         if parent_node:
            if action == QtCore.Qt.CopyAction:
               print "copy action, but we are moving"
               new_node = cur_node
            elif action == QtCore.Qt.MoveAction:
               print "move action, but we copy"
               new_node = osg2.cloneTree(cur_node)
            
            if -1 == row:
               parent_node.addChild(new_node)
            else:
               assert (row <= parent_node.getNChildren())
               parent_node.insertChild(row,new_node)
            self.emit(QtCore.SIGNAL("dataChanged(QModelIndex, QModelIndex)"), parent, parent)         
            self.emit(QtCore.SIGNAL("layoutChanged()"))         
         
      else:
         return False
   
      return True
      

   def insertRows(self, row, count, parent):
      print "insertRows: (%s, %s, %s)"%(row,count,parent)
      self.beginInsertRows(QtCore.QModelIndex(), row, row + count - 1)
      #for i in xrange(count):
      #   new_element = ET.SubElement(self.mElement, "cluster_node", name="NewNode", hostname="NewNode")
      #   new_node = ClusterNode(new_element)
      #   self.mNodes.insert(row, new_node);
      #self.refreshConnections()
      self.endInsertRows()
      #self.emit(QtCore.SIGNAL("rowsInserted(int, int)"), row, count)
      return True

   def removeRows(self, row, count, parent):
      print "removeRows: (%s, %s, %s)"%(row,count,parent)
      self.beginRemoveRows(QtCore.QModelIndex(), row, row + count - 1)
      #for i in xrange(count):
      #   node = self.mNodes[row]
      #
      #   # Remove node's element from XML tree.
      #   self.mElement.remove(node.mElement)
      #   # Remove node data structure
      #   self.mNodes.remove(node)
      self.endRemoveRows()
      return True

   # --- Change interface --- #
   def onNeedLayoutChanged(self):
      """ Called when we need to tell views to refresh their view. """
      print "emitting layout changed."
      self.emit(QtCore.SIGNAL("layoutChanged()"))

